package com.asgard.crypto;

import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Arrays;

import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.engines.AESFastEngine;
import org.bouncycastle.crypto.generators.SCrypt;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

/**
 * 该类封装了对密钥进行加密的函数。
 */
@SuppressWarnings("deprecation")
public class KeyCrypter {

    /**
     * 是否使用慢哈希
     */
    private static final boolean USE_SLOW_HASH = true;
    private static final SecureRandom secureRandom;

    static {
        secureRandom = new SecureRandom();
        Security.addProvider(new BouncyCastleProvider());
    }

    /**
     * 加密参数SALT_LENGTH
     */
    public static final int SALT_LENGTH = 8;
    /**
     * 加密参数KEY_LENGTH
     */
    public static final int KEY_LENGTH = 32;
    /**
     * 加密参数BLOCK_LENGTH
     */
    public static final int BLOCK_LENGTH = 16;
    /**
     * 加密参数ITERATIONS
     */
    public static final long ITERATIONS = 16384;

    /**
     * 初始化加密的SALT。
     *
     * @return 加密的SALT
     */
    public static byte[] randomSalt() {
        byte[] salt = new byte[SALT_LENGTH];
        secureRandom.nextBytes(salt);
        return salt;
    }

    /**
     * 用随机的SALT和ITERATIONS初始化。
     */
    public KeyCrypter() {
        this(randomSalt(), ITERATIONS);
    }

    /**
     * 用指定的SALT初始化。
     *
     * @param salt SALT
     */
    public KeyCrypter(byte[] salt) {
        this(salt, ITERATIONS);
    }

    /**
     * 用指定的ITERATIONS初始化。
     *
     * @param N ITERATIONS
     */
    public KeyCrypter(long N) {
        this(randomSalt(), N);
    }

    private final byte[] salt;
    private final long N;

    /**
     * 用指定的SALT和ITERATIONS初始化。
     *
     * @param salt SALT
     * @param N    ITERATIONS
     */
    public KeyCrypter(byte[] salt, long N) {
        this.salt = salt;
        this.N = N;
    }

    /**
     * 从密码生成KeyData。
     *
     * @param password 密码
     * @return KeyData
     */
    public KeyData deriveKey(char[] password) {
        if (!USE_SLOW_HASH)
            return new KeyData(CryptoUtils.encodeSha256(convertToByteArray(password)));
        byte[] passwordBytes = null;
        try {
            passwordBytes = convertToByteArray(password);
            byte[] keyBytes = SCrypt.generate(passwordBytes, salt, (int) N, 8, 1, KEY_LENGTH);
            return new KeyData(keyBytes);
        } catch (Exception e) {
            throw new RuntimeException("Could not generate key from password and salt.", e);
        } finally {
            if (passwordBytes != null)
                Arrays.fill(passwordBytes, (byte) 0);
        }
    }

    /**
     * 用密码加密。
     *
     * @param data   需要加密的数据
     * @param aesKey KeyData
     * @return 加密后的数据
     */
    public EncryptedData encrypt(byte[] data, KeyData aesKey) {
        try {
            byte[] iv = new byte[BLOCK_LENGTH];
            secureRandom.nextBytes(iv);
            BufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(new AESFastEngine()));
            cipher.init(true, new ParametersWithIV(new KeyParameter(aesKey.key), iv));
            byte[] encryptedBytes = new byte[cipher.getOutputSize(data.length)];
            final int length1 = cipher.processBytes(data, 0, data.length, encryptedBytes, 0);
            final int length2 = cipher.doFinal(encryptedBytes, length1);
            return new EncryptedData(iv, Arrays.copyOf(encryptedBytes, length1 + length2));
        } catch (Exception e) {
            throw new RuntimeException("Could not encrypt bytes.", e);
        }
    }

    /**
     * 用密码解密。
     *
     * @param encryptedData 加密后的数据
     * @param aesKey        KeyData
     * @return 解密后的数据
     */
    public byte[] decrypt(EncryptedData encryptedData, KeyData aesKey) {
        try {
            ParametersWithIV keyWithIv = new ParametersWithIV(new KeyParameter(aesKey.key),
                    encryptedData.initialisationVector);
            BufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(new AESFastEngine()));
            cipher.init(false, keyWithIv);
            byte[] cipherBytes = encryptedData.encryptedBytes;
            byte[] decryptedBytes = new byte[cipher.getOutputSize(cipherBytes.length)];
            final int length1 = cipher.processBytes(cipherBytes, 0, cipherBytes.length, decryptedBytes, 0);
            final int length2 = cipher.doFinal(decryptedBytes, length1);
            return Arrays.copyOf(decryptedBytes, length1 + length2);
        } catch (Exception e) {
            throw new RuntimeException("Could not decrypt bytes", e);
        }
    }

    /**
     * 将字符串转换成为字节数组
     *
     * @param chars 字符数组表示的密码
     * @return 字节数组
     */
    private static byte[] convertToByteArray(char[] chars) {
        CharBuffer charBuffer = CharBuffer.wrap(chars);
        ByteBuffer byteBuffer = StandardCharsets.UTF_8.encode(charBuffer);
        byte[] bytes = Arrays.copyOfRange(byteBuffer.array(), byteBuffer.position(), byteBuffer.limit());
        // 清除内存中的敏感数据
        Arrays.fill(charBuffer.array(), '\u0000');
        Arrays.fill(byteBuffer.array(), (byte) 0);
        return bytes;
    }

}
